#!/bin/bash

set -e

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 默认配置
SNAPSHOTTER_NAME="smart"
TEST_IMAGE="docker.io/library/nginx:latest"
CONTAINER_PREFIX="concurrent-test"
MAX_CONCURRENT=10
TIMEOUT=300

# 日志文件
LOG_DIR="./test-logs"
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
LOG_FILE="$LOG_DIR/concurrent-test-$TIMESTAMP.log"

# 性能数据文件
PERF_DATA="$LOG_DIR/performance-$TIMESTAMP.json"

# 创建日志目录
mkdir -p "$LOG_DIR"

# 记录日志函数
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

error() {
    echo -e "${RED}[ERROR] $1${NC}" | tee -a "$LOG_FILE"
}

success() {
    echo -e "${GREEN}[SUCCESS] $1${NC}" | tee -a "$LOG_FILE"
}

warning() {
    echo -e "${YELLOW}[WARNING] $1${NC}" | tee -a "$LOG_FILE"
}

# 检查依赖
check_dependencies() {
    log "检查依赖项..."

    if ! command -v ctr &> /dev/null; then
        error "ctr 命令未找到，请安装 containerd"
        exit 1
    fi

    if ! command -v jq &> /dev/null; then
        error "jq 命令未找到，请安装 jq"
        exit 1
    fi

    success "依赖项检查通过"
}

# 检查快照器是否可用
check_snapshotter() {
    log "检查快照器 $SNAPSHOTTER_NAME 是否可用..."

    if ! ctr plugin ls | grep -q "$SNAPSHOTTER_NAME"; then
        error "快照器 $SNAPSHOTTER_NAME 未找到"
        exit 1
    fi

    success "快照器 $SNAPSHOTTER_NAME 可用"
}

# 拉取测试镜像
pull_test_image() {
    log "拉取测试镜像 $TEST_IMAGE..."

    if ! ctr image ls | grep -q "nginx:latest"; then
        if timeout 300 ctr image pull "$TEST_IMAGE"; then
            success "测试镜像拉取成功"
        else
            error "测试镜像拉取失败"
            exit 1
        fi
    else
        log "测试镜像已存在"
    fi
}

# 清理测试容器
cleanup_containers() {
    log "清理测试容器..."

    local containers=$(ctr container ls | grep "$CONTAINER_PREFIX" | awk '{print $1}' || true)

    if [ -n "$containers" ]; then
        echo "$containers" | xargs -r ctr container rm -f || true
        log "已清理现有测试容器"
    fi
}

# 启动单个容器并记录性能数据
start_container() {
    local id=$1
    local container_name="${CONTAINER_PREFIX}-${id}"
    local start_time=$(date +%s.%N)

    log "启动容器 $container_name (ID: $id)..."

    # 记录内存使用
    local mem_before=$(free -m | awk 'NR==2{printf "%.2f", $3*100/$2}')

    # 启动容器并测试nginx配置
    if timeout "$TIMEOUT" ctr run --snapshotter "$SNAPSHOTTER_NAME" --rm "$TEST_IMAGE" "$container_name" nginx -t > /dev/null 2>&1; then
        local end_time=$(date +%s.%N)
        local duration=$(echo "$end_time - $start_time" | bc -l)
        local mem_after=$(free -m | awk 'NR==2{printf "%.2f", $3*100/$2}')

        # 获取容器信息
        local snapshot_info=$(ctr container info "$container_name" 2>/dev/null | jq -r '.Snapshotter' || echo "unknown")

        success "容器 $container_name 启动成功 (${duration}s)"

        # 记录性能数据
        echo "{\"id\": $id, \"duration\": $duration, \"mem_before\": $mem_before, \"mem_after\": $mem_after, \"snapshotter\": \"$snapshot_info\"}" >> "$PERF_DATA"

        return 0
    else
        local end_time=$(date +%s.%N)
        local duration=$(echo "$end_time - $start_time" | bc -l)

        error "容器 $container_name 启动失败 (${duration}s)"

        # 记录失败数据
        echo "{\"id\": $id, \"duration\": $duration, \"status\": \"failed\"}" >> "$PERF_DATA"

        return 1
    fi
}

# 并发启动测试
concurrent_start_test() {
    log "开始并发启动测试 (最大并发数: $MAX_CONCURRENT)..."

    local pids=()
    local success_count=0
    local failure_count=0
    local total_start_time=$(date +%s.%N)

    # 清空性能数据文件
    echo "[" > "$PERF_DATA"

    # 并发启动容器
    for i in $(seq 1 "$MAX_CONCURRENT"); do
        start_container "$i" &
        pids+=($!)

        # 限制并发数，避免系统过载
        if [ $((i % 5)) -eq 0 ]; then
            log "等待前5个容器启动完成..."
            for pid in "${pids[@]: -5}"; do
                wait "$pid" || ((failure_count++))
            done
        fi
    done

    # 等待所有容器启动完成
    log "等待所有容器启动完成..."
    for pid in "${pids[@]}"; do
        if wait "$pid"; then
            ((success_count++))
        else
            ((failure_count++))
        fi
    done

    local total_end_time=$(date +%s.%N)
    local total_duration=$(echo "$total_end_time - $total_start_time" | bc -l)

    # 完成性能数据文件
    sed -i '$ s/,$//' "$PERF_DATA"  # 移除最后一个逗号
    echo "]" >> "$PERF_DATA"

    # 生成测试报告
    generate_report "$success_count" "$failure_count" "$total_duration"
}

# 生成测试报告
generate_report() {
    local success_count=$1
    local failure_count=$2
    local total_duration=$3

    log "生成测试报告..."

    local report_file="$LOG_DIR/concurrent-test-report-$TIMESTAMP.txt"

    cat > "$report_file" << EOF
=====================================
并发容器启动测试报告
=====================================
测试时间: $(date)
快照器: $SNAPSHOTTER_NAME
测试镜像: $TEST_IMAGE
并发数量: $MAX_CONCURRENT
总耗时: ${total_duration}秒

结果统计:
- 成功: $success_count
- 失败: $failure_count
- 成功率: $(echo "scale=2; $success_count * 100 / ($success_count + $failure_count)" | bc -l)%

详细性能数据: $PERF_DATA
完整日志: $LOG_FILE

性能分析:
EOF

    # 如果有性能数据，进行分析
    if [ -f "$PERF_DATA" ] && [ -s "$PERF_DATA" ]; then
        local avg_duration=$(jq -r '[.[].duration] | add / length' "$PERF_DATA" 2>/dev/null || echo "N/A")
        local max_duration=$(jq -r '[.[].duration] | max' "$PERF_DATA" 2>/dev/null || echo "N/A")
        local min_duration=$(jq -r '[.[].duration] | min' "$PERF_DATA" 2>/dev/null || echo "N/A")

        cat >> "$report_file" << EOF
- 平均启动时间: ${avg_duration}秒
- 最长启动时间: ${max_duration}秒
- 最短启动时间: ${min_duration}秒
EOF
    fi

    success "测试报告已生成: $report_file"

    # 显示报告内容
    echo
    cat "$report_file"
}

# 系统资源监控
monitor_resources() {
    log "系统资源使用情况:"

    # CPU 使用率
    local cpu_usage=$(top -bn1 | grep "Cpu(s)" | sed "s/.*, *\([0-9.]*\)%* id.*/\1/" | awk '{print 100 - $1}')
    echo "CPU 使用率: ${cpu_usage}%" | tee -a "$LOG_FILE"

    # 内存使用率
    local mem_usage=$(free | grep Mem | awk '{printf "%.2f", $3/$2 * 100.0}')
    echo "内存使用率: ${mem_usage}%" | tee -a "$LOG_FILE"

    # 磁盘使用率
    local disk_usage=$(df -h / | awk 'NR==2 {print $5}')
    echo "磁盘使用率: $disk_usage" | tee -a "$LOG_FILE"

    # 系统负载
    local load_avg=$(uptime | awk -F'load average:' '{print $2}')
    echo "系统负载:$load_avg" | tee -a "$LOG_FILE"
}

# 主测试函数
main() {
    log "开始并发容器启动测试"
    log "====================================="

    # 检查依赖
    check_dependencies

    # 检查快照器
    check_snapshotter

    # 拉取测试镜像
    pull_test_image

    # 清理测试容器
    cleanup_containers

    # 显示初始资源状态
    monitor_resources

    # 执行并发测试
    concurrent_start_test

    # 显示最终资源状态
    monitor_resources

    # 最终清理
    cleanup_containers

    log "并发容器启动测试完成"
    log "====================================="
}

# 参数解析
while [[ $# -gt 0 ]]; do
    case $1 in
        --snapshotter)
            SNAPSHOTTER_NAME="$2"
            shift 2
            ;;
        --image)
            TEST_IMAGE="$2"
            shift 2
            ;;
        --concurrent)
            MAX_CONCURRENT="$2"
            shift 2
            ;;
        --timeout)
            TIMEOUT="$2"
            shift 2
            ;;
        --help)
            echo "用法: $0 [选项]"
            echo "选项:"
            echo "  --snapshotter NAME    快照器名称 (默认: smart)"
            echo "  --image IMAGE        测试镜像 (默认: docker.io/library/nginx:latest)"
            echo "  --concurrent NUM     并发数量 (默认: 10)"
            echo "  --timeout SECONDS    超时时间 (默认: 300)"
            echo "  --help               显示帮助信息"
            exit 0
            ;;
        *)
            error "未知参数: $1"
            exit 1
            ;;
    esac
done

# 运行主函数
main "$@"